import os
import numpy as np
import torch
import torch.nn as nn
from torch.utils.data import DataLoader
import pandas as pd
from tqdm.auto import tqdm
from model_define import FoodCNN
from train_transform import FoodDataset, train_tfm, test_tfm
from torch.utils.tensorboard import SummaryWriter
from model_vgg import ModularVGG11
writer = SummaryWriter('logs')

def calculate_class_weights(train_dataset):
    class_counts = np.zeros(11)
    for _, label in train_dataset:
        class_counts[label] += 1
    class_weights = 1.0 / (class_counts + 1e-6)
    class_weights = class_weights / class_weights.sum() * 11
    return torch.tensor(class_weights, dtype=torch.float)

def train_model():
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    train_dataset = FoodDataset('food11/versions/1/training', 'train', train_tfm)
    val_dataset = FoodDataset('food11/versions/1/validation', 'val', test_tfm)

    class_weights = calculate_class_weights(train_dataset).to(device)

    train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True, num_workers=2)
    val_loader = DataLoader(val_dataset, batch_size=32, shuffle=False, num_workers=2)

    model = ModularVGG11(num_classes=11).to(device)

    criterion = nn.CrossEntropyLoss(weight=class_weights)
    optimizer = torch.optim.Adam(model.parameters(), lr=0.001, weight_decay=1e-5)
    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', patience=5, factor=0.5)
    # 训练循环
    best_acc = 0.0
    for epoch in range(50):
        model.train()
        train_loss = 0.0
        train_correct = 0
        train_total = 0

        pbar = tqdm(train_loader, desc=f'Epoch {epoch + 1}/50')
        for images, labels in pbar:
            images, labels = images.to(device), labels.to(device)
            # 前向传播
            outputs = model(images)
            loss = criterion(outputs, labels)
            # 反向传播
            optimizer.zero_grad()
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=10)
            optimizer.step()
            # 统计信息
            train_loss += loss.item()
            _, predicted = torch.max(outputs, 1)
            train_total += labels.size(0)
            train_correct += (predicted == labels).sum().item()
            pbar.set_postfix({'Loss': f'{loss.item():.4f}', 'Acc': f'{train_correct / train_total:.4f}'})
        model.eval()
        val_loss = 0.0
        val_correct = 0
        val_total = 0
        with torch.no_grad():
            for images, labels in val_loader:
                images, labels = images.to(device), labels.to(device)
                outputs = model(images)
                loss = criterion(outputs, labels)

                val_loss += loss.item()
                _, predicted = torch.max(outputs, 1)
                val_total += labels.size(0)
                val_correct += (predicted == labels).sum().item()

        train_acc = 100.0 * train_correct / train_total
        val_acc = 100.0 * val_correct / val_total

        print(f'Epoch {epoch + 1}: Train Loss: {train_loss / len(train_loader):.4f}, '
              f'Train Acc: {train_acc:.2f}%, Val Loss: {val_loss / len(val_loader):.4f}, '
              f'Val Acc: {val_acc:.2f}%')
        writer.add_scalar('loss', val_loss, epoch)
        writer.add_scalar('accuracy', val_acc, epoch)

        scheduler.step(val_loss / len(val_loader))


        if val_acc > best_acc:
            best_acc = val_acc
            torch.save(model.state_dict(), 'best_food_model_vgg.pth')
            print(f'保存最佳模型，验证集准确率: {val_acc:.2f}%')
    print(f'训练完成，最佳验证集准确率: {best_acc:.2f}%')
    return model

def predict(model, test_loader, device):
    model.eval()
    predictions = []
    filenames = []

    with torch.no_grad():
        for images, names in tqdm(test_loader, desc='预测中'):
            images = images.to(device)
            outputs = model(images)
            _, predicted = torch.max(outputs, 1)

            predictions.extend(predicted.cpu().numpy())
            filenames.extend(names)

    results = []
    for i, (fn, pred) in enumerate(zip(filenames, predictions)):
        img_id = os.path.splitext(fn)[0]  # 去除.jpg后缀
        results.append({'Id': img_id, 'Category': int(pred)})

    result_df = pd.DataFrame(results)
    result_df.to_csv('ans_ours.csv', index=False)
    print("预测结果已保存到 ans_ours.csv")

    return result_df

# 主执行函数
def main():
    torch.manual_seed(42)
    np.random.seed(42)
    # 训练模型
    print("开始训练模型...")
    trained_model = train_model()
    print("加载最佳模型进行预测...")
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    best_model = FoodCNN(num_classes=11)
    best_model.load_state_dict(torch.load('best_food_model_vgg.pth'))
    best_model.to(device)

    test_dataset = FoodDataset('food11/versions/1/evaluation', 'test', test_tfm)
    test_loader = DataLoader(test_dataset, batch_size=128, shuffle=False, num_workers=2)

    # 预测
    predictions = predict(best_model, test_loader, device)
    print("预测完成！")
    print(predictions.head())

    print("\n预测结果类别分布:")
    print(predictions['Category'].value_counts().sort_index())

if __name__ == "__main__":
    main()